{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "b33f31a9-cc41-4132-85c6-ccb6e5ca0d41",
   "metadata": {},
   "outputs": [],
   "source": [
    "from model.Conv_TasNet.hand.conv_tasnet import ConvTasNet as Model\n",
    "import torch\n",
    "import torchaudio"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 56,
   "id": "618db3ed-c150-4537-8140-70786fa80fbb",
   "metadata": {},
   "outputs": [],
   "source": [
    "checkpoint = torch.load(\"convTasNetV1.pth\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 57,
   "id": "84dbb2a4-324a-4bfa-b821-c420aef7b717",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "<All keys matched successfully>"
      ]
     },
     "execution_count": 57,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m = Model(512, sr=8000).to(\"cpu\")\n",
    "m.load_state_dict(checkpoint['model_state_dict'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 58,
   "id": "3129e29d-f825-4401-aa9a-36b5ff2fa0ac",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "ConvTasNet(\n",
       "  (encoder): Conv1d(1, 512, kernel_size=(16,), stride=(8,), bias=False)\n",
       "  (seperator): TCN(\n",
       "    (layerNorm): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "    (BN): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "    (stacks): ModuleList(\n",
       "      (0): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(1,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (1): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(2,), dilation=(2,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (2): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(4,), dilation=(4,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (3): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(8,), dilation=(8,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (4): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(16,), dilation=(16,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (5): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(32,), dilation=(32,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (6): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(64,), dilation=(64,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (7): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(128,), dilation=(128,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (8): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(1,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (9): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(2,), dilation=(2,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (10): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(4,), dilation=(4,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (11): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(8,), dilation=(8,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (12): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(16,), dilation=(16,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (13): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(32,), dilation=(32,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (14): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(64,), dilation=(64,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (15): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(128,), dilation=(128,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (16): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(1,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (17): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(2,), dilation=(2,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (18): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(4,), dilation=(4,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (19): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(8,), dilation=(8,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (20): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(16,), dilation=(16,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (21): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(32,), dilation=(32,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (22): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(64,), dilation=(64,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "      (23): Conv1DBlock(\n",
       "        (conv1d): Conv1d(128, 512, kernel_size=(1,), stride=(1,))\n",
       "        (dconv): Conv1d(512, 512, kernel_size=(3,), stride=(1,), padding=(128,), dilation=(128,), groups=512)\n",
       "        (res_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "        (pReLU1): PReLU(num_parameters=1)\n",
       "        (pReLU2): PReLU(num_parameters=1)\n",
       "        (norm1): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (norm2): GroupNorm(1, 512, eps=1e-08, affine=True)\n",
       "        (skip_out): Conv1d(512, 128, kernel_size=(1,), stride=(1,))\n",
       "      )\n",
       "    )\n",
       "    (output): Sequential(\n",
       "      (0): PReLU(num_parameters=1)\n",
       "      (1): Conv1d(128, 1024, kernel_size=(1,), stride=(1,))\n",
       "    )\n",
       "  )\n",
       "  (decoder): ConvTranspose1d(512, 1, kernel_size=(16,), stride=(8,), bias=False)\n",
       ")"
      ]
     },
     "execution_count": 58,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "m.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 59,
   "id": "dd9983ae-638d-45c9-a80e-08d8aedc6109",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[ 0.0000,  0.0000,  0.0000,  ..., -0.0003, -0.0002, -0.0004]]), 8000)"
      ]
     },
     "execution_count": 59,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "torchaudio.load(\"mix_n.wav\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 73,
   "id": "eecfd940-1f60-40f8-8073-2ee96b7201d8",
   "metadata": {},
   "outputs": [],
   "source": [
    "# y = torchaudio.load(\"1.wav\")[0]\n",
    "# y1 = torchaudio.load(\"1.wav\")[0]\n",
    "x = torchaudio.load(\"mix.wav\")[0]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 74,
   "id": "3c3c941d-f320-4841-86a4-774416c914a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "y = m(x)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 76,
   "id": "6844a11f-a39d-4baa-93a4-da8669a814af",
   "metadata": {},
   "outputs": [],
   "source": [
    "torchaudio.save(\"o1.wav\",y[0][1].unsqueeze(0), 8000)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "09506963",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f1851ce2",
   "metadata": {},
   "outputs": [],
   "source": [
    "x1 = torch.randn(2,1,2,5)\n",
    "x2 = torch.randn(2,2,1,5)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
